import logging
import re
import os
import json

from agent.common.system import sysCommand
from agent.common.config import Config
from agent.feature.base import Feature
from agent.plugin.common import read_file

'''
Feature: govergor
Domain: cpu
'''

logger = logging.getLogger('common')

class Governor(Feature):
    __domain__ = "cpu"
    __feature__ = "governor"
    
    def __init__(self) -> None:
        """ Check available of this feature

        Check rpm file, script file and service availability
        Raise an exception to indicate that the feature is not available
        e.g.
            raise Exception("lib file do not exists")
            raise Exception("feature available on ARM only")
                
        Raises:
            Exception: feature is not available
        """
        super().__init__(self.__domain__, self.__feature__)
        try:
            sysCommand("which cpupower")
        except Exception as e:
            msg = "disable governor control: {}".format(e)
            logger.warning(msg)
            raise Exception(msg)
        
        governor_info = sysCommand("cpupower -c all frequency-info",log=False)
        cpu_governor_info_list = [i for i in re.split(r"analyzing CPU \d+:", governor_info) if i!=""]
        self.cpu_avaliable_governor = {}
        for index, cpu_governor_info in enumerate(cpu_governor_info_list):
            avaliable_governor_match = re.search(r"available cpufreq governors:(.*)",cpu_governor_info)
            if avaliable_governor_match is not None:
                avaliable_governor = [i.strip() for i in avaliable_governor_match.group(1).split(' ') if i!='' \
                                and i in ['performance', 'powersave', 'ondemand', 'userspace', 'conservative']]
                if avaliable_governor.__len__() != 0:
                    self.cpu_avaliable_governor[index] = avaliable_governor
                logger.debug("CPU{index} available cpufreq governors:{avaliable_governor}".format(
                    index = index,
                    avaliable_governor = avaliable_governor,
                ))

    # Tuned is_cpu_online()
    def _is_cpu_online(self, cpu_index):
		# CPU0 is always online
        if cpu_index == 0:
            return True
        
        if read_file("/sys/devices/system/cpu/cpu{}/online".format(cpu_index)).strip()=="1":
            return True
        else:
            return False

    # Tuned _cpu_has_scaling_governor()
    def _cpu_has_scaling_governor(self, cpu_index):
        return os.path.exists("/sys/devices/system/cpu/cpu{}/cpufreq/scaling_governor".format(cpu_index))

    # Tuned _check_cpu_can_change_governor()
    def _check_cpu_can_change_governor(self, cpu_index):
        if not self._is_cpu_online(cpu_index):
            logger.info("cpu{} is not online, skipping".format(cpu_index))
            return False
        if not self._cpu_has_scaling_governor(cpu_index):
            logger.debug("there is no scaling governor for cpu{}, skipping".format(cpu_index))
            return False
        return True

    # Tuned _set_governor()
    def _parseValidGovernor(self, governor, avaliable_list):
        """ select valid governor
        
        e.g. ondemand|powersave -> powersave, if powersave in avaliable_list
        """
        governor_list = governor.split('|')
        for _gov in governor_list:
            if _gov in avaliable_list:
                return _gov

    def set_value(self, value):
        """ Setting parameter value
        Args:
            value (string): param value

        Raises:
            Exception: Setting value failed
        """
        for cpu_index in self.cpu_avaliable_governor.keys():
            if not self._check_cpu_can_change_governor(cpu_index):
                continue
            
            valid_governor = self._parseValidGovernor(
                value, self.cpu_avaliable_governor[cpu_index])
            
            if valid_governor is None:
                logger.warning("invalid governor {} for cpu{}".format(value, cpu_index))
            else:
                sysCommand("cpupower -c {index} frequency-set -g {param_value}".format(
                    index = cpu_index,
                    param_value = valid_governor))

    def get_value(self):
        """ Reading parameter value in env

        Reading parameter values in the current environment.

        Returns:
            string: parameter value

        Raises:
            Exception: Reading value failed
        """
        govergor = {}
        for cpu_index in self.cpu_avaliable_governor.keys():
            if not self._check_cpu_can_change_governor(cpu_index):
                continue

            data = read_file("/sys/devices/system/cpu/cpu{}/cpufreq/scaling_governor".format(cpu_index)).strip()
            if len(data) > 0:
                govergor[cpu_index] = data
            
            else:
                logger.error("could not get current governor on cpu{}".format(cpu_index))
                raise Exception("could not get current governor on cpu{}".format(cpu_index))
        return govergor
    
    # rewrite rollcak() function
    def rollback(self):
        backup_file = os.path.join(
            Config.BACKUP_PATH, "{}_{}.json".format(self.domain_name, self.feature_name))
        if not os.path.exists(backup_file):
            return
        
        logger.info("load backup data from: {}".format(backup_file))
        value = json.load(open(backup_file))
        for cpu_index, govergor in value.items():
            if not self._check_cpu_can_change_governor(cpu_index):
                continue
            sysCommand("cpupower -c {index} frequency-set -g {govergor}".format(
                    index = cpu_index,
                    govergor = govergor))
        
        os.remove(backup_file)

if __name__ == "__main__":
    # Unit Test
    try:
        handler = Governor()
        handler.get_value()
        handler.backup()
        handler.set_value("performance")
        handler.set_value("conservative|powersave")
        handler.rollback()
    except Exception as e:
        print(e)